import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style({"font.sans-serif": ['Microsoft YaHei', 'SimHei']})  # 显示中文

rn3 = pd.read_csv('../data2.csv', index_col='Date')  # 读取数据
rn3_rets = np.log(rn3.pct_change() + 1)  # 计算对数收益率,这里的pct_change是严格的收益率，加了对数以便和正太分布对上
rn3_rets.index = [str(x)[5:7] for x in rn3_rets.index]  # 将Date转换为月份
rn3_group = rn3_rets.groupby(rn3_rets.index).sum()  # 数据聚合，得到每个月的总收益率
rn3_group.index.name = '月份'
rn3_group.columns.name = '股票'
plt.figure(figsize=(10, 8))
# 其实这个函数可以理解成将这个dataframe显示出来而已
sns.heatmap(rn3_group, annot=True, linewidths=.5)
plt.show()
# sns.pairplot(rn3_rets.dropna())
